#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;


class KernelReduceSum_fp32 {
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX; //inQueueT保存每次tiling的中间结果，需要EnQue DeQue等操作
    // TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<TPosition::VECCALC>calcBuf, yBuf; //reduce所需要的临时数据

    GlobalTensor<float> xGm;
    GlobalTensor<float> yGm;

    uint32_t minLength;
    uint32_t ALIGN_NUM;

    uint32_t* shape;
    uint32_t dim;

    int32_t axesNum;

    uint32_t blockLength;
    uint32_t blockSize;
    uint32_t unitCount;
    uint32_t totalLength; //总数据个数
    uint32_t stepSize;
    
    uint32_t tileNum; //每组reduce数据计算的循环次数
    uint32_t tileLength; //每次tiling的数据长度
    uint32_t lasttileLength; //最后一次tiling的长度/ uint32_t remain_size;

public:
    __aicore__  inline KernelReduceSum_fp32() {}
    __aicore__  inline  void Init(GM_ADDR x, GM_ADDR axes, GM_ADDR y, uint32_t blockLength, uint32_t blockSize, uint32_t ALIGN_NUM, uint32_t* shape, uint32_t dim)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->blockLength = blockLength;
        this->blockSize = blockSize;
        this->ALIGN_NUM = ALIGN_NUM;
        this->minLength = this->ALIGN_NUM;
        this->shape = shape;
        this->dim = dim;

        this->axesNum = (*(__gm__ int32_t *)axes);
        if(axesNum < 0)
        {
            axesNum = dim + axesNum;
        }
        
        this->totalLength = shape[axesNum];
        
        this->stepSize = 1;
        this->unitCount = 1;
        for(uint32_t i = 0; i < axesNum; i++)
        {
            this->unitCount *= shape[i];
        }

        for(uint32_t i = axesNum+1; i < dim; i++)
        {
            this->stepSize *= shape[i];
        }

        uint32_t totalLengthAligned = (totalLength + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

        this->tileNum = totalLengthAligned / this->blockSize;
        // 
        if(totalLengthAligned % blockSize == 0 || tileNum == 0) //reduce dim对齐之后，可以均分到各个tile，或不足一个tile
        {
            if(tileNum == 0)    //一次tiling的数据个数 大于 totalLengthAligned，也就是tile_num = 0时，tileLength=totalLengthAligned
            {
                this->tileNum = 1;
                this->tileLength = totalLengthAligned;
                this->lasttileLength = this->tileLength;
            }
            else //一次tiling的数据个数 小于 totalLengthAligned时，tileLength取最大值tiling_size * ALIGN_NUM
            {
                this->tileLength = this->blockSize;
                this->lasttileLength = this->tileLength;
            }
            

        }
        else //totalLengthAligned对齐后不能均分到每次tiling
        {
            this->tileNum = this->tileNum + 1;
            this->tileLength = this->blockSize; //tileLength 取最大值(ub)
            this->lasttileLength = totalLengthAligned % this->tileLength;
        }

        // printf("tileNum : %u \n", tileNum);
        // printf("tileLength : %u \n", tileLength);
        // printf("lasttileLength : %u \n", lasttileLength);
        // printf("unitCount : %u \n", unitCount);
        // printf("stepSize : %u \n", stepSize);
        // printf("totalLength : %u \n", totalLength);

        ASSERT(tileNum != 0 && "tile num can not be zero!");
    
        
        xGm.SetGlobalBuffer((__gm__ float *)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ float *)y, unitCount * stepSize);
        
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(float));
        // pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(float)); //存储float的结果
        pipe.InitBuffer(calcBuf, BUFFER_NUM * this->tileLength * sizeof(float));
        pipe.InitBuffer(yBuf, BUFFER_NUM * this->minLength * sizeof(float));    
    }

    __aicore__ inline void Process()
    {
        if (stepSize != 1)
        {
            // 处理轴 != -1情况, 需要setvalue来设置xlocal
            ProcessAxes();
            return;
        }
        // 处理 最后一维度做sum的情况,可使用datacopy优化
        uint32_t loopCount = this->tileNum;

        for (int32_t k = 0; k < this->unitCount; k++)
        {
            float sum = 0.0;
            for(int32_t i = 0; i < loopCount - 1; i++)
            {
                CopyIn(i,  k, tileLength);
                Compute(i, tileLength, sum);
            }
            CopyIn(tileNum - 1,  k, lasttileLength);
            Compute(tileNum - 1, totalLength - (tileNum - 1) * tileLength, sum);

            auto cursor = yGm[k].address_;
            *cursor = sum;
        }
    }


    __aicore__ inline void CopyIn(int32_t progress, int k, uint32_t length)
    {
        LocalTensor<float> xLocal = inQueueX.AllocTensor<float>();
        uint32_t startIndex = k  * this->totalLength + progress * tileLength;
        DataCopy(xLocal, xGm[startIndex], length);
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void ProcessAxes()
    {
        uint32_t loopCount = this->tileNum;
        for (int32_t k = 0; k < this->unitCount; k++)
        {
            for (int32_t j = 0; j < this->stepSize; j++)
            {
                float sum = 0.0;
                for(int32_t i = 0; i < loopCount - 1 ; i++)
                {
                    CopyInAxes(i, j, k, tileLength);
                    Compute(i, tileLength, sum);
                }
                CopyInAxes(tileNum - 1, j, k, totalLength - (tileNum -1)*tileLength);
                Compute(tileNum - 1, totalLength - (tileNum -1)*tileLength, sum);

                auto cursor = yGm[j + k * this->stepSize].address_;
                *cursor = sum;
            }   
        }
    }

    
    __aicore__ inline void CopyInAxes(int32_t progress, int j, int k, uint32_t length)
    {
        LocalTensor<float> xLocal = inQueueX.AllocTensor<float>();
        uint32_t startIndex = k * this->stepSize * this->totalLength + progress * tileLength * this->stepSize + j;

        for (int i = 0; i < length; i++)
        {
            xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
        }
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress, uint32_t length, float& sum)
    {   
        LocalTensor<float> xLocal = inQueueX.DeQue<float>();
        LocalTensor<float> yLocal = yBuf.Get<float>();
        LocalTensor<float> wLocal = calcBuf.Get<float>();

        ReduceSum(xLocal, xLocal, wLocal, length);
        Adds(yLocal, xLocal, (float)sum, minLength);
        sum = yLocal.GetValue(0);
        inQueueX.FreeTensor(xLocal);
    }
};

class KernelReduceSum_fp16 {
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX; //inQueueT保存每次tiling的中间结果，需要EnQue DeQue等操作
    // TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<TPosition::VECCALC>calcBuf, xBuf, yBuf; //reduce所需要的临时数据

    GlobalTensor<half> xGm;
    GlobalTensor<half> yGm;
    
    uint32_t minLength;
    uint32_t ALIGN_NUM;

    uint32_t* shape;
    uint32_t dim;

    int32_t axesNum;

    uint32_t blockLength;
    uint32_t blockSize;
    uint32_t unitCount;
    uint32_t totalLength; //总数据个数
    uint32_t stepSize;
    
    uint32_t tileNum; //每组reduce数据计算的循环次数
    uint32_t tileLength; //每次tiling的数据长度
    uint32_t lasttileLength; //最后一次tiling的长度/ uint32_t remain_size;

public:
    __aicore__  inline KernelReduceSum_fp16() {}
    __aicore__  inline  void Init(GM_ADDR x, GM_ADDR axes, GM_ADDR y, uint32_t blockLength, uint32_t blockSize, uint32_t ALIGN_NUM, uint32_t* shape, uint32_t dim)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");


        this->blockLength = blockLength;
        this->blockSize = blockSize;
        this->ALIGN_NUM = ALIGN_NUM;
        this->minLength = this->ALIGN_NUM;
        this->shape = shape;
        this->dim = dim;

        this->axesNum = (*(__gm__ int32_t *)axes);
        if(axesNum < 0)
        {
            axesNum = dim + axesNum;
        }
        // printf("axesNum : %d \n", axesNum);
        
        this->totalLength = shape[axesNum];
        
        this->stepSize = 1;
        this->unitCount = 1;
        for(uint32_t i = 0; i < axesNum; i++)
        {
            this->unitCount *= shape[i];
        }

        for(uint32_t i = axesNum+1; i < dim; i++)
        {
            this->stepSize *= shape[i];
        }

        uint32_t totalLengthAligned = (totalLength + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

        this->tileNum = totalLengthAligned / this->blockSize;
        // 
        if(totalLengthAligned % blockSize == 0 || tileNum == 0) //reduce dim对齐之后，可以均分到各个tile，或不足一个tile
        {
            if(tileNum == 0)    //一次tiling的数据个数 大于 totalLengthAligned，也就是tile_num = 0时，tileLength=totalLengthAligned
            {
                this->tileNum = 1;
                this->tileLength = totalLengthAligned;
                this->lasttileLength = this->tileLength;
            }
            else //一次tiling的数据个数 小于 totalLengthAligned时，tileLength取最大值tiling_size * ALIGN_NUM
            {
                this->tileLength = this->blockSize;
                this->lasttileLength = this->tileLength;
            }
        }
        else //totalLengthAligned对齐后不能均分到每次tiling
        {
            this->tileNum = this->tileNum + 1;
            this->tileLength = this->blockSize; //tileLength 取最大值(ub)
            this->lasttileLength = totalLengthAligned % this->tileLength;
        }

        // printf("tileNum : %u \n", tileNum);
        // printf("tileLength : %u \n", tileLength);
        // printf("lasttileLength : %u \n", lasttileLength);
        // printf("unitCount : %u \n", unitCount);
        // printf("stepSize : %u \n", stepSize);
        // printf("totalLength : %u \n", totalLength);
        ASSERT(tileNum != 0 && "tile num can not be zero!");
    
        xGm.SetGlobalBuffer((__gm__ half*)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ half*)y, unitCount * stepSize);
        
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(half));
        // pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(half)); //存储float的结果
        pipe.InitBuffer(calcBuf, BUFFER_NUM * this->tileLength * sizeof(float));
        pipe.InitBuffer(xBuf, BUFFER_NUM * this->tileLength * sizeof(float));    
        pipe.InitBuffer(yBuf, BUFFER_NUM * this->minLength * sizeof(float));    
    }

    __aicore__ inline void Process()
    {
        if (stepSize != 1)
        {
            // 处理轴 != -1情况, 需要setvalue来设置xlocal
            ProcessAxes();
            return;
        }
        // 处理 最后一维度做sum的情况,可使用datacopy优化
        uint32_t loopCount = this->tileNum;

        for (int32_t k = 0; k < this->unitCount; k++)
        {
            float sum = 0.0;
            for(int32_t i = 0; i < loopCount - 1; i++)
            {
                CopyIn(i,  k, tileLength);
                Compute(i, tileLength, sum);
            }
            CopyIn(tileNum - 1,  k, lasttileLength);
            Compute(tileNum - 1, totalLength - (tileNum - 1) * tileLength, sum);
            
            auto cursor = yGm[k].address_;
            *cursor = (half)(sum);
        }
    }


    __aicore__ inline void CopyIn(int32_t progress, int k, uint32_t length)
    {
        LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
        uint32_t startIndex = k  * this->totalLength + progress * tileLength;
        DataCopy(xLocal, xGm[startIndex], length);
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void ProcessAxes()
    {
        uint32_t loopCount = this->tileNum;
        for (int32_t k = 0; k < this->unitCount; k++)
        {
            for (int32_t j = 0; j < this->stepSize; j++)
            {
                float sum = 0.0;
                for(int32_t i = 0; i < loopCount - 1; i++)
                {
                    CopyInAxes(i, j, k, tileLength);
                    Compute(i, tileLength, sum);
                }
                CopyInAxes(tileNum -1, j, k, totalLength - (tileNum -1) * tileLength);
                Compute(tileNum -1, totalLength - (tileNum -1) * tileLength, sum);

                auto cursor = yGm[j + k * this->stepSize].address_;
                *cursor = (half)(sum);
            }   
        }
    }

    
    __aicore__ inline void CopyInAxes(int32_t progress, int j, int k, uint32_t length)
    {
        LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
        uint32_t startIndex = k * stepSize * totalLength + progress * tileLength * stepSize + j;
        for (int i = 0; i < length; i++)
        {
            xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
        }
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress, uint32_t length, float& sum)
    {   
        LocalTensor<half> xLocal = inQueueX.DeQue<half>();
        LocalTensor<float> yLocal = yBuf.Get<float>();
        LocalTensor<float> xLocal_f = xBuf.Get<float>();
        LocalTensor<float> wLocal = calcBuf.Get<float>();

        Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, length);
        ReduceSum(xLocal_f, xLocal_f, wLocal, length);
        Adds(yLocal, xLocal_f, sum, minLength);
        sum = yLocal.GetValue(0);
        inQueueX.FreeTensor(xLocal);
    }
};

class KernelReduceSum_int32 {
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX; //inQueueT保存每次tiling的中间结果，需要EnQue DeQue等操作
    // TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<TPosition::VECCALC>calcBuf, xBuf, yBuf; //reduce所需要的临时数据

    GlobalTensor<int32_t> xGm;
    GlobalTensor<int32_t> yGm;

    uint32_t minLength;
    uint32_t ALIGN_NUM;

    uint32_t* shape;
    uint32_t dim;

    int32_t axesNum;

    uint32_t blockLength;
    uint32_t blockSize;
    uint32_t unitCount;
    uint32_t totalLength; //总数据个数
    uint32_t stepSize;
    
    uint32_t tileNum; //每组reduce数据计算的循环次数
    uint32_t tileLength; //每次tiling的数据长度
    uint32_t lasttileLength; //最后一次tiling的长度/ uint32_t remain_size;

public:
    __aicore__  inline KernelReduceSum_int32() {}
    __aicore__  inline  void Init(GM_ADDR x, GM_ADDR axes, GM_ADDR y, uint32_t blockLength, uint32_t blockSize, uint32_t ALIGN_NUM, uint32_t* shape, uint32_t dim)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->blockLength = blockLength;
        this->blockSize = blockSize;
        this->ALIGN_NUM = ALIGN_NUM;
        this->minLength = this->ALIGN_NUM;
        this->shape = shape;
        this->dim = dim;

        this->axesNum = (*(__gm__ int32_t *)axes);
        if(axesNum < 0)
        {
            axesNum = dim + axesNum;
        }
        
        this->totalLength = shape[axesNum];
        
        this->stepSize = 1;
        this->unitCount = 1;
        for(uint32_t i = 0; i < axesNum; i++)
        {
            this->unitCount *= shape[i];
        }

        for(uint32_t i = axesNum+1; i < dim; i++)
        {
            this->stepSize *= shape[i];
        }

        uint32_t totalLengthAligned = (totalLength + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

        this->tileNum = totalLengthAligned / this->blockSize;
        // 
        if(totalLengthAligned % blockSize == 0 || tileNum == 0) //reduce dim对齐之后，可以均分到各个tile，或不足一个tile
        {
            if(tileNum == 0)    //一次tiling的数据个数 大于 totalLengthAligned，也就是tile_num = 0时，tileLength=totalLengthAligned
            {
                this->tileNum = 1;
                this->tileLength = totalLengthAligned;
                this->lasttileLength = this->tileLength;
            }
            else //一次tiling的数据个数 小于 totalLengthAligned时，tileLength取最大值tiling_size * ALIGN_NUM
            {
                this->tileLength = this->blockSize;
                this->lasttileLength = this->tileLength;
            }
            

        }
        else //totalLengthAligned对齐后不能均分到每次tiling
        {
            this->tileNum = this->tileNum + 1;
            this->tileLength = this->blockSize; //tileLength 取最大值(ub)
            this->lasttileLength = totalLengthAligned % this->tileLength;
        }

        // printf("tileNum : %u \n", tileNum);
        // printf("tileLength : %u \n", tileLength);
        // printf("lasttileLength : %u \n", lasttileLength);
        // printf("unitCount : %u \n", unitCount);
        // printf("stepSize : %u \n", stepSize);
        // printf("totalLength : %u \n", totalLength);

        ASSERT(tileNum != 0 && "tile num can not be zero!");
    
        
        xGm.SetGlobalBuffer((__gm__ int32_t *)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ int32_t *)y, unitCount * stepSize);
        
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(int32_t));
        // pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(float)); //存储float的结果
        pipe.InitBuffer(calcBuf, BUFFER_NUM * this->tileLength * sizeof(float));
        pipe.InitBuffer(xBuf, BUFFER_NUM * this->tileLength * sizeof(float));    
        pipe.InitBuffer(yBuf, BUFFER_NUM * this->minLength * sizeof(float));    
    }

    __aicore__ inline void Process()
    {
        if (stepSize != 1)
        {
            // 处理轴 != -1情况, 需要setvalue来设置xlocal
            ProcessAxes();
            return;
        }
        // 处理 最后一维度做sum的情况,可使用datacopy优化
        uint32_t loopCount = this->tileNum;

        for (int32_t k = 0; k < this->unitCount; k++)
        {
            float sum = 0.0;
            for(int32_t i = 0; i < loopCount - 1; i++)
            {
                CopyIn(i,  k, tileLength);
                Compute(i, tileLength, sum);
            }
            CopyIn(tileNum - 1,  k, lasttileLength);
            Compute(tileNum - 1, totalLength - (tileNum - 1) * tileLength, sum);

            auto cursor = yGm[k].address_;
            *cursor = (int32_t)sum;
        }
    }


    __aicore__ inline void CopyIn(int32_t progress, int k, uint32_t length)
    {
        LocalTensor<int32_t> xLocal = inQueueX.AllocTensor<int32_t>();
        uint32_t startIndex = k  * this->totalLength + progress * tileLength;
        DataCopy(xLocal, xGm[startIndex], length);
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void ProcessAxes()
    {
        uint32_t loopCount = this->tileNum;
        for (int32_t k = 0; k < this->unitCount; k++)
        {
            for (int32_t j = 0; j < this->stepSize; j++)
            {
                float sum = 0.0;
                for(int32_t i = 0; i < loopCount - 1; i++)
                {
                    CopyInAxes(i, j, k, tileLength);
                    Compute(i, tileLength, sum);
                }
                CopyInAxes(tileNum - 1, j, k, totalLength - (tileNum - 1) * tileLength);
                Compute(tileNum - 1, totalLength - (tileNum - 1) * tileLength, sum);

                auto cursor = yGm[j + k * this->stepSize].address_;
                *cursor = (int32_t)sum;
            }
        }
    }

    
    __aicore__ inline void CopyInAxes(int32_t progress, int j, int k, uint32_t length)
    {
        LocalTensor<int32_t> xLocal = inQueueX.AllocTensor<int32_t>();
        uint32_t startIndex = k * this->stepSize * this->totalLength + progress * tileLength * this->stepSize + j;

        for (int i = 0; i < length; i++)
        {
            xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
        }
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress, uint32_t length, float& sum)
    {   
        LocalTensor<int32_t> xLocal = inQueueX.DeQue<int32_t>();
        LocalTensor<float> yLocal = yBuf.Get<float>();
        LocalTensor<float> xLocal_f = xBuf.Get<float>();
        LocalTensor<float> wLocal = calcBuf.Get<float>();

        Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, length);

        ReduceSum(xLocal_f, xLocal_f, wLocal, length);
        Adds(yLocal, xLocal_f, (float)sum, minLength);
        sum = yLocal.GetValue(0);
        inQueueX.FreeTensor(xLocal);
    }
};

class KernelReduceSum_int8 {
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX; //inQueueT保存每次tiling的中间结果，需要EnQue DeQue等操作
    // TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<TPosition::VECCALC>calcBuf, xBuf, yBuf; //reduce所需要的临时数据

    GlobalTensor<int8_t> xGm;
    GlobalTensor<int8_t> yGm;

    uint32_t minLength;
    uint32_t ALIGN_NUM;

    uint32_t* shape;
    uint32_t dim;

    int32_t axesNum;

    uint32_t blockLength;
    uint32_t blockSize;
    uint32_t unitCount;
    uint32_t totalLength; //总数据个数
    uint32_t stepSize;
    
    uint32_t tileNum; //每组reduce数据计算的循环次数
    uint32_t tileLength; //每次tiling的数据长度
    uint32_t lasttileLength; //最后一次tiling的长度/ uint32_t remain_size;

public:
    __aicore__  inline KernelReduceSum_int8() {}
    __aicore__  inline  void Init(GM_ADDR x, GM_ADDR axes, GM_ADDR y, uint32_t blockLength, uint32_t blockSize, uint32_t ALIGN_NUM, uint32_t* shape, uint32_t dim)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->blockLength = blockLength;
        this->blockSize = blockSize;
        this->ALIGN_NUM = ALIGN_NUM;
        this->minLength = this->ALIGN_NUM;
        this->shape = shape;
        this->dim = dim;

        this->axesNum = (*(__gm__ int32_t *)axes);
        if(axesNum < 0)
        {
            axesNum = dim + axesNum;
        }
        
        this->totalLength = shape[axesNum];
        
        this->stepSize = 1;
        this->unitCount = 1;
        for(uint32_t i = 0; i < axesNum; i++)
        {
            this->unitCount *= shape[i];
        }

        for(uint32_t i = axesNum+1; i < dim; i++)
        {
            this->stepSize *= shape[i];
        }

        uint32_t totalLengthAligned = (totalLength + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

        this->tileNum = totalLengthAligned / this->blockSize;
        // 
        if(totalLengthAligned % blockSize == 0 || tileNum == 0) //reduce dim对齐之后，可以均分到各个tile，或不足一个tile
        {
            if(tileNum == 0)    //一次tiling的数据个数 大于 totalLengthAligned，也就是tile_num = 0时，tileLength=totalLengthAligned
            {
                this->tileNum = 1;
                this->tileLength = totalLengthAligned;
                this->lasttileLength = this->tileLength;
            }
            else //一次tiling的数据个数 小于 totalLengthAligned时，tileLength取最大值tiling_size * ALIGN_NUM
            {
                this->tileLength = this->blockSize;
                this->lasttileLength = this->tileLength;
            }
        }
        else //totalLengthAligned对齐后不能均分到每次tiling
        {
            this->tileNum = this->tileNum + 1;
            this->tileLength = this->blockSize; //tileLength 取最大值(ub)
            this->lasttileLength = totalLengthAligned % this->tileLength;
        }

        // printf("tileNum : %u \n", tileNum);
        // printf("tileLength : %u \n", tileLength);
        // printf("lasttileLength : %u \n", lasttileLength);
        // printf("unitCount : %u \n", unitCount);
        // printf("stepSize : %u \n", stepSize);
        // printf("totalLength : %u \n", totalLength);

        ASSERT(tileNum != 0 && "tile num can not be zero!");
    
        
        xGm.SetGlobalBuffer((__gm__ int8_t *)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ int8_t *)y, unitCount * stepSize);
        
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(int8_t));
        // pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(float)); //存储float的结果
        pipe.InitBuffer(calcBuf, BUFFER_NUM * this->tileLength * sizeof(half));
        pipe.InitBuffer(xBuf, BUFFER_NUM * this->tileLength * sizeof(half));
        pipe.InitBuffer(yBuf, BUFFER_NUM * this->minLength * sizeof(half));
    }

    __aicore__ inline void Process()
    {
        if (stepSize != 1)
        {
            // 处理轴 != -1情况, 需要setvalue来设置xlocal
            ProcessAxes();
            return;
        }
        // 处理 最后一维度做sum的情况,可使用datacopy优化
        uint32_t loopCount = this->tileNum;

        for (int32_t k = 0; k < this->unitCount; k++)
        {
            half sum = 0.0;
            for(int32_t i = 0; i < loopCount - 1; i++)
            {
                CopyIn(i,  k, tileLength);
                Compute(i, tileLength, sum);
            }
            CopyIn(tileNum - 1,  k, lasttileLength);
            Compute(tileNum - 1, totalLength - (tileNum - 1) * tileLength, sum);

            auto cursor = yGm[k].address_;
            *cursor = (int8_t)sum;
        }
    }


    __aicore__ inline void CopyIn(int32_t progress, int k, uint32_t length)
    {
        LocalTensor<int8_t> xLocal = inQueueX.AllocTensor<int8_t>();
        uint32_t startIndex = k  * this->totalLength + progress * tileLength;
        DataCopy(xLocal, xGm[startIndex], length);
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void ProcessAxes()
    {
        uint32_t loopCount = this->tileNum;
        for (int32_t k = 0; k < this->unitCount; k++)
        {
            for (int32_t j = 0; j < this->stepSize; j++)
            {
                half sum = 0.0;
                for(int32_t i = 0; i < loopCount - 1; i++)
                {
                    CopyInAxes(i, j, k, tileLength);
                    Compute(i, tileLength, sum);
                }
                CopyInAxes(tileNum -1, j, k, totalLength - (tileNum -1) * tileLength);
                Compute(tileNum -1, totalLength - (tileNum -1) * tileLength, sum);

                auto cursor = yGm[j + k * this->stepSize].address_;
                *cursor = (int8_t)sum;
            }
        }
    }

    
    __aicore__ inline void CopyInAxes(int32_t progress, int j, int k, uint32_t length)
    {
        LocalTensor<int8_t> xLocal = inQueueX.AllocTensor<int8_t>();
        uint32_t startIndex = k * this->stepSize * this->totalLength + progress * tileLength * this->stepSize + j;

        for (int i = 0; i < length; i++)
        {
            xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
        }
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress, uint32_t length, half& sum)
    {   
        LocalTensor<int8_t> xLocal = inQueueX.DeQue<int8_t>();
        LocalTensor<half> yLocal = yBuf.Get<half>();
        LocalTensor<half> xLocal_f = xBuf.Get<half>();
        LocalTensor<half> wLocal = calcBuf.Get<half>();

        Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, length);

        ReduceSum(xLocal_f, xLocal_f, wLocal, length);
        Adds(yLocal, xLocal_f, sum, minLength);
        sum = yLocal.GetValue(0);
        inQueueX.FreeTensor(xLocal);
    }
};



extern "C" __global__ __aicore__ void reduce_sum(GM_ADDR x, GM_ADDR axes, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    if(TILING_KEY_IS(1))
    {
        KernelReduceSum_fp32 op;
        op.Init(x, axes, y, tiling_data.blockLength, tiling_data.blockSize, tiling_data.ALIGN_NUM, tiling_data.shape, tiling_data.dimNum);
        op.Process();
    }
    else if (TILING_KEY_IS(2))
    {
        KernelReduceSum_fp16 op;
        op.Init(x, axes, y, tiling_data.blockLength, tiling_data.blockSize, tiling_data.ALIGN_NUM, tiling_data.shape, tiling_data.dimNum);
        op.Process();
    }
    else if (TILING_KEY_IS(3))
    {
        KernelReduceSum_int32 op;
        op.Init(x, axes, y, tiling_data.blockLength, tiling_data.blockSize, tiling_data.ALIGN_NUM, tiling_data.shape, tiling_data.dimNum);
        op.Process();
    }
    else if (TILING_KEY_IS(4))
    {
        KernelReduceSum_int8 op;
        op.Init(x, axes, y, tiling_data.blockLength, tiling_data.blockSize, tiling_data.ALIGN_NUM, tiling_data.shape, tiling_data.dimNum);
        op.Process();
    }
}